import cv2
import numpy as np
from setuptools import glob
import os


# raw_pic_dir 存放原始数据的文件夹，以*.png结尾
# npz结果输出的文件夹
def npz(raw_pic_dir, npz_output_dir, num_classes=2):
    datanames = os.listdir(raw_pic_dir.split("*")[0])
    print("the num_classes is %d", num_classes)
    # path = glob.escape(path)
    # path2 = glob.escape(path2)
    for i, img_path in enumerate(glob.glob(raw_pic_dir)):
        # 读入图像
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 读入标签
        label_path = img_path.replace('raw', 'affined')
        label = cv2.imread(label_path, flags=0)
        if num_classes == 2:
            # 将非目标像素设置为0
            label[label != 255] = 0
            # 将目标像素设置为1
            label[label == 255] = 1

        else:  # todo 此处需要根据自己的需要更改分类的像素标签
            label[label != 255] = 0
            label[label == 255] = 1
        # 保存npz
        np.savez(npz_output_dir + str(datanames[i].split(".")[0]), image=image, label=label)
        print(datanames[i].split(".")[0])
        print('------------', i)


def write_name(npz_output_dir):

    npz_dir = npz_output_dir + r'\*.npz'
    files = glob.glob(npz_dir)
    # txt文件路径
    f = open(r'lists\lists_Synapse\test_vol.txt', 'w')
    for i in files:
        name = i.split('\\')[-1]
        name = name[:-4] + '\n'
        f.write(name)
    # npz文件路径


if "__main__":
    npz()
    write_name()
    print('ok')
